import os
import cv2
import numpy as np
from PIL import Image, ImageFilter
import uuid
from datetime import datetime
import traceback

# 为OpenCV函数添加类型注释
from typing import Any, Tuple

class IDPhotoProcessor:
    def __init__(self):
        # 创建处理结果存储目录
        self.output_dir = "processed_photos"
        if not os.path.exists(self.output_dir):
            os.makedirs(self.output_dir)
            
        # 定义标准尺寸规格 (像素)
        self.size_specs = {
            '1x': (295, 413),    # 一寸
            '2x': (413, 579),    # 二寸
            's1x': (260, 378),   # 小一寸
            's2x': (413, 531)    # 小二寸
        }
        
        # 定义背景颜色BGR值
        self.background_colors = {
            'white': (255, 255, 255),
            'blue': (180, 100, 60),   # 蓝色背景 (BGR)
            'red': (60, 60, 200)      # 红色背景 (BGR)
        }
    
    def process_photo(self, image_path, settings):
        """
        处理证件照
        :param image_path: 原始图片路径
        :param settings: 处理设置
        :return: 处理后的图片路径和信息
        """
        try:
            # 检查文件是否存在
            if not os.path.exists(image_path):
                raise Exception(f"图片文件不存在: {image_path}")
            
            # 读取图片
            image = cv2.imread(image_path)  # type: ignore
            if image is None:
                raise Exception("无法读取图片，请确保图片格式正确")
            
            # 调整尺寸
            resized_image = self._resize_image(image, settings)
            
            # 更换背景
            processed_image = self._change_background(resized_image, settings)
            
            # 去水印（如果启用）
            if settings.get('removeWatermark', False):
                processed_image = self._remove_watermark(processed_image)
            
            # 保存处理后的图片
            output_path = self._save_image(processed_image)
            
            # 获取文件大小
            file_size = self._get_file_size(output_path)
            
            return {
                'processed_path': output_path,
                'file_size': file_size
            }
            
        except Exception as e:
            # 打印详细的错误信息
            error_info = traceback.format_exc()
            print(f"处理图片时出错: {e}")
            print(f"详细错误信息: {error_info}")
            raise Exception(f"处理图片时出错: {str(e)}")
    
    def _resize_image(self, image: Any, settings: dict) -> Any:
        """
        调整图片尺寸
        """
        try:
            size_type = settings.get('size', '1x')
            
            if size_type == 'custom':
                width = int(settings.get('customWidth', 295))
                height = int(settings.get('customHeight', 413))
            else:
                width, height = self.size_specs.get(size_type, (295, 413))
            
            # 确保尺寸为正数
            width = max(1, width)
            height = max(1, height)
            
            # 调整图片尺寸
            resized = cv2.resize(image, (width, height), interpolation=cv2.INTER_AREA)  # type: ignore
            return resized
        except Exception as e:
            raise Exception(f"调整图片尺寸时出错: {str(e)}")
    
    def _change_background(self, image: Any, settings: dict) -> Any:
        """
        更换背景颜色（使用更精确的背景替换方法）
        """
        try:
            background_type = settings.get('background', 'white')
            
            if background_type == 'custom':
                # 使用自定义颜色
                color_hex = settings.get('customColor', '#FFFFFF')
                # 将十六进制颜色转换为BGR
                color_bgr = self._hex_to_bgr(color_hex)
            else:
                # 使用预定义颜色
                color_bgr = self.background_colors.get(background_type, (255, 255, 255))
            
            # 创建纯色背景
            background = np.full(image.shape, color_bgr, dtype=np.uint8)
            
            # 创建人物掩码
            mask = self._create_human_mask_advanced(image)
            
            # 对掩码进行优化处理
            mask = self._refine_mask(mask)
            
            # 添加羽化效果使边缘更自然
            mask = self._feather_mask(mask, 3)
            
            # 将掩码转换为3通道
            mask_3channel = cv2.cvtColor(mask, cv2.COLOR_GRAY2BGR)  # type: ignore
            
            # 归一化掩码到0-1范围
            mask_normalized = mask_3channel.astype(float) / 255
            
            # 使用掩码混合原图和背景
            result = (image * mask_normalized + background * (1 - mask_normalized)).astype(np.uint8)
            
            return result
        except Exception as e:
            # 如果背景替换失败，使用备用方法
            return self._change_background_fallback(image, settings)
    
    def _create_human_mask_advanced(self, image: Any) -> Any:
        """
        创建人物掩码（使用更高级的背景分离方法）
        """
        try:
            # 1. 基于肤色的掩码
            skin_mask = self._improved_skin_detection(image)
            
            # 2. 基于边缘的掩码
            edge_mask = self._edge_detection(image)
            
            # 3. 基于颜色分布的掩码
            color_mask = self._color_distribution_mask(image)
            
            # 4. 基于人脸检测的掩码
            face_mask = self._face_detection_mask(image)
            
            # 5. 基于连通区域的掩码
            connected_mask = self._connected_components_mask(skin_mask)
            
            # 6. 结合多种方法的结果
            # 使用加权平均来结合不同的掩码
            combined_mask = cv2.addWeighted(skin_mask, 0.25, edge_mask, 0.15, 0)  # type: ignore
            combined_mask = cv2.addWeighted(combined_mask, 1.0, color_mask, 0.2, 0)  # type: ignore
            combined_mask = cv2.addWeighted(combined_mask, 1.0, face_mask, 0.2, 0)  # type: ignore
            combined_mask = cv2.addWeighted(combined_mask, 1.0, connected_mask, 0.2, 0)  # type: ignore
            
            # 7. 使用形态学操作优化掩码
            kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))  # type: ignore
            combined_mask = cv2.morphologyEx(combined_mask, cv2.MORPH_CLOSE, kernel)  # type: ignore
            combined_mask = cv2.morphologyEx(combined_mask, cv2.MORPH_OPEN, kernel)  # type: ignore
            
            # 8. 使用更精确的GrabCut进行最终优化
            final_mask = self._improved_grabcut_refinement(image, combined_mask)
            
            return final_mask
        except Exception as e:
            # 如果高级方法失败，返回基础肤色检测结果
            return self._improved_skin_detection(image)
    
    def _change_background_fallback(self, image: Any, settings: dict) -> Any:
        """备用的背景替换方法"""
        try:
            background_type = settings.get('background', 'white')
            if background_type == 'custom':
                color_hex = settings.get('customColor', '#FFFFFF')
                color_bgr = self._hex_to_bgr(color_hex)
            else:
                color_bgr = self.background_colors.get(background_type, (255, 255, 255))

            background = np.full(image.shape, color_bgr, dtype=np.uint8)
            result = cv2.addWeighted(image, 0.8, background, 0.2, 0)  # type: ignore
            return result
        except:
            # 如果所有方法都失败，返回原图
            return image
    def _face_detection(self, image: Any) -> Any:
        """基于人脸检测创建掩码"""
        try:
            # 加载人脸检测器
            face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml')  # type: ignore

            # 转换为灰度图
            gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)  # type: ignore

            # 检测人脸
            faces = face_cascade.detectMultiScale(gray, 1.1, 4)

            # 创建人脸掩码
            face_mask = np.zeros(image.shape[:2], dtype=np.uint8)

            # 在检测到的人脸区域绘制椭圆
            for (x, y, w, h) in faces:
                # 在人脸区域绘制椭圆
                center = (x + w//2, y + h//2)
                axes = (w//2, h//2)
                cv2.ellipse(face_mask, center, axes, 0, 0, 360, 255, -1)  # type: ignore

            return face_mask
        except:
            # 如果人脸检测失败，返回空掩码
            return np.zeros(image.shape[:2], dtype=np.uint8)
    
    def _connected_components_mask(self, mask: Any) -> Any:
        """
        基于连通区域分析创建掩码
        """
        try:
            # 查找连通区域
            num_labels, labels, stats, centroids = cv2.connectedComponentsWithStats(mask, connectivity=8)  # type: ignore
            
            # 创建新的掩码，只保留最大的几个连通区域
            refined_mask = np.zeros_like(mask)
            
            # 获取区域面积，排除背景（标签0）
            areas = stats[1:, cv2.CC_STAT_AREA]  # type: ignore
            
            # 如果找到了连通区域
            if len(areas) > 0:
                # 找到面积最大的几个区域（最多5个）
                num_regions = min(5, len(areas))
                largest_indices = np.argpartition(areas, -num_regions)[-num_regions:]
                
                # 在新掩码中绘制这些区域
                for idx in largest_indices:
                    # 注意：stats中的索引需要+1，因为排除了背景
                    label = idx + 1
                    refined_mask[labels == label] = 255
            
            return refined_mask
        except:
            return mask
    
    def _improved_grabcut_refinement(self, image: Any, initial_mask: Any) -> Any:
        """
        改进的GrabCut算法进行精细的前景提取
        """
        try:
            # 创建GrabCut所需的数组
            mask = np.zeros(image.shape[:2], np.uint8)
            bgdModel = np.zeros((1, 65), np.float64)
            fgdModel = np.zeros((1, 65), np.float64)
            
            # 创建确定的背景和前景区域
            h, w = image.shape[:2]
            
            # 设置确定的背景区域（图像边缘）
            border_width = max(5, min(h, w) // 20)
            mask[:border_width, :] = cv2.GC_BGD  # type: ignore
            mask[-border_width:, :] = cv2.GC_BGD  # type: ignore
            mask[:, :border_width] = cv2.GC_BGD  # type: ignore
            mask[:, -border_width:] = cv2.GC_BGD  # type: ignore
            
            # 基于初始掩码设置可能的前景区域
            mask[(initial_mask > 50) & (initial_mask <= 200)] = cv2.GC_PR_FGD  # type: ignore
            mask[initial_mask > 200] = cv2.GC_FGD  # type: ignore
            mask[initial_mask <= 50] = cv2.GC_PR_BGD  # type: ignore
            
            # 应用GrabCut算法
            cv2.grabCut(image, mask, None, bgdModel, fgdModel, 10, cv2.GC_INIT_WITH_MASK)  # type: ignore
            
            # 提取前景掩码
            mask2 = np.where((mask == cv2.GC_FGD) | (mask == cv2.GC_PR_FGD), 255, 0).astype('uint8')  # type: ignore
            
            # 对结果进行后处理
            kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3))  # type: ignore
            mask2 = cv2.morphologyEx(mask2, cv2.MORPH_CLOSE, kernel)  # type: ignore
            mask2 = cv2.morphologyEx(mask2, cv2.MORPH_OPEN, kernel)  # type: ignore
            
            return mask2
        except:
            # 如果GrabCut失败，返回初始掩码
            return initial_mask
    
    def _improved_skin_detection(self, image: Any) -> Any:
        """
        改进的肤色检测算法
        """
        try:
            # 转换为不同的色彩空间
            hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)  # type: ignore
            ycrcb = cv2.cvtColor(image, cv2.COLOR_BGR2YCrCb)  # type: ignore
            lab = cv2.cvtColor(image, cv2.COLOR_BGR2LAB)  # type: ignore
            
            # HSV色彩空间的肤色范围
            lower_hsv = np.array([0, 15, 0], dtype=np.uint8)
            upper_hsv = np.array([17, 170, 255], dtype=np.uint8)
            mask_hsv = cv2.inRange(hsv, lower_hsv, upper_hsv)  # type: ignore
            
            # YCrCb色彩空间的肤色范围
            lower_ycrcb = np.array([0, 135, 85], dtype=np.uint8)
            upper_ycrcb = np.array([255, 180, 135], dtype=np.uint8)
            mask_ycrcb = cv2.inRange(ycrcb, lower_ycrcb, upper_ycrcb)  # type: ignore
            
            # Lab色彩空间的肤色范围
            lower_lab = np.array([0, 125, 125], dtype=np.uint8)
            upper_lab = np.array([255, 140, 140], dtype=np.uint8)
            mask_lab = cv2.inRange(lab, lower_lab, upper_lab)  # type: ignore
            
            # 结合不同色彩空间的结果
            combined_mask = cv2.bitwise_and(mask_hsv, mask_ycrcb)  # type: ignore
            combined_mask = cv2.bitwise_and(combined_mask, mask_lab)  # type: ignore
            
            # 应用高斯模糊减少噪声
            combined_mask = cv2.GaussianBlur(combined_mask, (3, 3), 0)  # type: ignore
            
            # 应用阈值确保二值化
            _, combined_mask = cv2.threshold(combined_mask, 50, 255, cv2.THRESH_BINARY)  # type: ignore
            
            return combined_mask
        except:
            return np.zeros(image.shape[:2], dtype=np.uint8)
    
    def _skin_detection(self, image: Any) -> Any:
        """
        基于肤色检测创建掩码
        """
        try:
            # 转换为HSV色彩空间
            hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)  # type: ignore
            
            # 定义更精确的肤色范围（HSV色彩空间）
            lower_skin = np.array([0, 20, 70], dtype=np.uint8)
            upper_skin = np.array([20, 255, 255], dtype=np.uint8)
            
            # 创建肤色掩码
            skin_mask_hsv = cv2.inRange(hsv, lower_skin, upper_skin)  # type: ignore
            
            # 转换为YCrCb色彩空间
            ycrcb = cv2.cvtColor(image, cv2.COLOR_BGR2YCrCb)  # type: ignore
            
            # YCrCb色彩空间的肤色范围
            lower_ycrcb = np.array([0, 135, 85], dtype=np.uint8)
            upper_ycrcb = np.array([255, 170, 125], dtype=np.uint8)
            
            # 创建YCrCb肤色掩码
            skin_mask_ycrcb = cv2.inRange(ycrcb, lower_ycrcb, upper_ycrcb)  # type: ignore
            
            # 合并两个掩码
            combined_skin_mask = cv2.bitwise_or(skin_mask_hsv, skin_mask_ycrcb)  # type: ignore
            
            return combined_skin_mask
        except:
            return np.zeros(image.shape[:2], dtype=np.uint8)
    
    def _edge_detection(self, image: Any) -> Any:
        """
        基于边缘检测创建掩码
        """
        try:
            # 转换为灰度图
            gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)  # type: ignore
            
            # 应用高斯模糊
            blurred = cv2.GaussianBlur(gray, (5, 5), 0)  # type: ignore
            
            # 使用Canny边缘检测
            edges = cv2.Canny(blurred, 50, 150)  # type: ignore
            
            # 使用形态学操作连接边缘
            kernel = np.ones((3, 3), np.uint8)
            edges = cv2.dilate(edges, kernel, iterations=1)  # type: ignore
            edges = cv2.erode(edges, kernel, iterations=1)  # type: ignore
            
            return edges
        except:
            return np.zeros(image.shape[:2], dtype=np.uint8)
    
    def _color_distribution_mask(self, image: Any) -> Any:
        """
        基于颜色分布创建掩码
        """
        try:
            # 计算图像的直方图
            hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)  # type: ignore
            
            # 计算2D直方图（H和S通道）
            hist = cv2.calcHist([hsv], [0, 1], None, [50, 60], [0, 180, 0, 256])  # type: ignore
            
            # 归一化直方图
            cv2.normalize(hist, hist, 0, 255, cv2.NORM_MINMAX)  # type: ignore
            
            # 反向投影
            back_proj = cv2.calcBackProject([hsv], [0, 1], hist, [0, 180, 0, 256], 1)  # type: ignore
            
            # 应用阈值
            _, mask = cv2.threshold(back_proj, 50, 255, cv2.THRESH_BINARY)  # type: ignore
            
            return mask
        except:
            return np.zeros(image.shape[:2], dtype=np.uint8)
    
    def _face_detection_mask(self, image: Any) -> Any:
        """
        基于人脸检测创建掩码
        """
        try:
            # 尝试加载人脸检测器
            face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml')  # type: ignore
            
            # 转换为灰度图
            gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)  # type: ignore
            
            # 检测人脸
            faces = face_cascade.detectMultiScale(gray, 1.1, 4)
            
            # 创建人脸掩码
            face_mask = np.zeros(image.shape[:2], dtype=np.uint8)
            
            # 在检测到的人脸区域绘制椭圆
            for (x, y, w, h) in faces:
                # 在人脸区域绘制椭圆
                center = (x + w//2, y + h//2)
                axes = (w//2, h//2)
                cv2.ellipse(face_mask, center, axes, 0, 0, 360, 255, -1)  # type: ignore
            
            return face_mask
        except:
            # 如果人脸检测失败，返回空掩码
            return np.zeros(image.shape[:2], dtype=np.uint8)
    
    def _grabcut_refinement(self, image: Any, initial_mask: Any) -> Any:
        """
        使用GrabCut算法进行精细的前景提取
        """
        try:
            # 创建GrabCut所需的数组
            mask = np.zeros(image.shape[:2], np.uint8)
            bgdModel = np.zeros((1, 65), np.float64)
            fgdModel = np.zeros((1, 65), np.float64)
            
            # 将初始掩码作为GrabCut的输入
            # 基于初始掩码设置GrabCut的初始状态
            mask[initial_mask == 0] = cv2.GC_BGD  # type: ignore
            mask[initial_mask == 255] = cv2.GC_FGD  # type: ignore
            mask[(initial_mask > 0) & (initial_mask < 255)] = cv2.GC_PR_FGD  # type: ignore
            
            # 如果没有明确的前景区域，使用矩形初始化
            if np.sum(mask == cv2.GC_FGD) == 0:  # type: ignore
                # 设置一个中心区域作为可能的前景
                h, w = image.shape[:2]
                rect = (w//4, h//4, w//2, h//2)
                cv2.grabCut(image, mask, rect, bgdModel, fgdModel, 5, cv2.GC_INIT_WITH_RECT)  # type: ignore
            else:
                # 使用掩码初始化
                cv2.grabCut(image, mask, None, bgdModel, fgdModel, 5, cv2.GC_INIT_WITH_MASK)  # type: ignore
            
            # 提取前景掩码
            mask2 = np.where((mask == cv2.GC_FGD) | (mask == cv2.GC_PR_FGD), 255, 0).astype('uint8')  # type: ignore
            
            return mask2
        except:
            # 如果GrabCut失败，返回初始掩码
            return initial_mask
    
    def _feather_mask(self, mask: Any, feather_radius: int) -> Any:
        """
        对掩码进行羽化处理，使边缘更自然
        """
        try:
            # 创建羽化核
            kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (feather_radius*2+1, feather_radius*2+1))  # type: ignore
            
            # 对掩码进行距离变换
            dist_transform = cv2.distanceTransform(mask, cv2.DIST_L2, 5)  # type: ignore
            
            # 归一化距离变换
            cv2.normalize(dist_transform, dist_transform, 0, 255, cv2.NORM_MINMAX)  # type: ignore
            
            # 创建羽化掩码
            feathered_mask = np.uint8(dist_transform)
            
            return feathered_mask
        except:
            return mask
    
    def _refine_mask(self, mask: Any) -> Any:
        """
        优化掩码质量
        """
        try:
            # 对掩码进行形态学操作以改善效果
            # 先进行开运算去除噪声
            kernel_open = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3))  # type: ignore
            mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel_open)  # type: ignore
            
            # 再进行闭运算填充空洞
            kernel_close = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))  # type: ignore
            mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel_close)  # type: ignore
            
            # 对掩码进行高斯模糊以使边缘更平滑
            mask = cv2.GaussianBlur(mask, (5, 5), 0)  # type: ignore
            
            # 使用自适应阈值确保掩码是二值的
            mask = cv2.adaptiveThreshold(mask, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY, 11, 2)  # type: ignore
            
            return mask
        except:
            # 如果自适应阈值失败，使用普通阈值
            try:
                _, mask = cv2.threshold(mask, 127, 255, cv2.THRESH_BINARY)  # type: ignore
                return mask
            except:
                return mask
    
    def _remove_watermark(self, image: Any) -> Any:
        """
        去水印处理（简化版本）
        """
        try:
            # 应用高斯模糊来减少水印可见性
            blurred = cv2.GaussianBlur(image, (3, 3), 0)  # type: ignore
            
            # 应用双边滤波器保持边缘清晰
            filtered = cv2.bilateralFilter(blurred, 9, 75, 75)  # type: ignore
            
            # 应用形态学操作
            kernel = np.ones((3,3), np.uint8)
            morphed = cv2.morphologyEx(filtered, cv2.MORPH_CLOSE, kernel)  # type: ignore
            
            return morphed
        except Exception as e:
            # 如果去水印失败，返回原图
            return image
    
    def _save_image(self, image: Any) -> str:
        """
        保存图片并返回路径
        """
        try:
            # 生成唯一文件名
            filename = f"processed_{uuid.uuid4().hex}_{int(datetime.now().timestamp())}.jpg"
            filepath = os.path.join(self.output_dir, filename)
            
            # 确保输出目录存在
            if not os.path.exists(self.output_dir):
                os.makedirs(self.output_dir)
            
            # 保存图片
            success = cv2.imwrite(filepath, image)  # type: ignore
            if not success:
                raise Exception("图片保存失败")
                
            return filepath
        except Exception as e:
            raise Exception(f"保存图片时出错: {str(e)}")
    
    def _get_file_size(self, filepath: str) -> str:
        """
        获取文件大小
        """
        try:
            size_bytes = os.path.getsize(filepath)
            if size_bytes < 1024:
                return f"{size_bytes} B"
            elif size_bytes < 1024 * 1024:
                return f"{size_bytes/1024:.1f} KB"
            else:
                return f"{size_bytes/(1024*1024):.1f} MB"
        except:
            return "未知大小"
    
    def _hex_to_bgr(self, hex_color: str) -> Tuple[int, int, int]:
        """
        将十六进制颜色转换为BGR
        """
        try:
            hex_color = hex_color.lstrip('#')
            # 确保hex_color有6个字符
            if len(hex_color) != 6:
                hex_color = 'FFFFFF'  # 默认白色
                
            rgb = tuple(int(hex_color[i:i+2], 16) for i in (0, 2, 4))
            return (rgb[2], rgb[1], rgb[0])  # 转换为BGR
        except:
            return (255, 255, 255)  # 默认白色

# 创建处理器实例
processor = IDPhotoProcessor()

def process_id_photo(image_path: str, settings: dict) -> dict:
    """
    处理证件照的外部接口
    """
    return processor.process_photo(image_path, settings)